Missing type annotations: The type annotations for training_image and algorithm_name were removed entirely. Per SDK coding standards (PEP 484), all public/private methods must retain type annotations. Since these parameters now accept both str and PipelineVariable, please use the appropriate union type:

def _validate_training_image_and_algorithm_name(
    self, training_image: str | PipelineVariable | None, algorithm_name: str | PipelineVariable | None
):

Or if StrPipeVar is already defined as a type alias in the codebase, use that.

Removed type annotation: Same issue here — the image parameter lost its type annotation. Please restore it with the correct union type:

def _get_repo_name_from_image(image: str | PipelineVariable) -> str:

Missing test file in the diff: The PR description references test_model_trainer_pipeline_variable.py but this file is not included in the changed files. Please ensure the test file is included in the PR. Without tests, we cannot verify the fix works or guard against regressions.

The mock session and mock_defaults setup is duplicated across 4 test methods (test_training_image_accepts_parameter_string, test_algorithm_name_accepts_parameter_string, test_environment_values_accept_parameter_string, test_plain_string_values_still_work). Extract this into a @pytest.fixture to reduce duplication and improve maintainability:

@pytest.fixture
def mock_session():
    session = MagicMock()
    session.boto_region_name = "us-east-1"
    session.default_bucket.return_value = "my-bucket"
    session.default_bucket_prefix = None
    return session

And similarly for the mock_defaults patching.

Consider adding a test case for _validate_training_image_and_algorithm_name where one argument is a PipelineVariable and the other is None — this is the primary success case the fix enables. The current tests test_training_image_accepts_parameter_string and test_algorithm_name_accepts_parameter_string test this indirectly through full ModelTrainer construction, but a direct unit test of the validation method (like the rejection tests at lines 199-238) would be more focused and faster.

Duplicate import: Union is imported here from typing, but there's already a from typing import Literal, Any on line 24. Consolidate into a single import statement:

from typing import Literal, Any, Union

Also, since the module already imports PipelineVariable from sagemaker.core.workflow.parameters on line 30, and from __future__ import annotations is not present, consider adding it to enable PEP 604 union syntax (str | PipelineVariable) per SDK conventions.

The logic isinstance(training_image, PipelineVariable) or bool(training_image) will raise TypeError if training_image is a PipelineVariable that doesn't support __bool__ — but since isinstance short-circuits via or, this is actually safe. However, consider simplifying to:

has_image = training_image is not None and training_image != ""
has_algo = algorithm_name is not None and algorithm_name != ""

This avoids calling bool() entirely and is more explicit about what "not provided" means (None or empty string). The is not None check naturally handles PipelineVariable objects correctly.

The return type annotation says -> str but when a PipelineVariable is passed, it returns a string placeholder, so the annotation is technically correct. However, consider documenting in the docstring that the placeholder _PIPELINE_VARIABLE_IMAGE_PLACEHOLDER is returned for PipelineVariable inputs, so downstream callers understand the behavior.

The from sagemaker.train.model_trainer import ModelTrainer import is repeated inside every test method in this class. Move it to the top of the file with the other imports. Inline imports in tests add unnecessary noise and are not consistent with SDK test conventions.

Nit: This test image URI contains a hardcoded region (us-east-1) and account ID (683313688378). While this is acceptable for unit tests since it's just a string constant and not used to make actual API calls, consider using a clearly fake account ID (e.g., 123456789012) for consistency with the mock session fixture below.

Using ModelTrainer.__new__(ModelTrainer) to bypass __init__ and directly test the validation method is fragile — it creates an uninitialized object. If _validate_training_image_and_algorithm_name ever accesses self attributes, these tests will break with confusing errors. Consider either:

1. Making _validate_training_image_and_algorithm_name a @staticmethod (it doesn't use self), or
2. Using the existing mock_train_defaults fixture to construct a proper instance and test through the public interface.

The TestSafeSerializeWithPipelineVariable tests verify safe_serialize behavior with PipelineVariable, but the PR diff doesn't show any changes to safe_serialize. If safe_serialize already handled PipelineVariable correctly, these tests are documenting existing behavior (which is fine), but it would be good to note that in the test class docstring. If safe_serialize needed changes, those changes should be included in this PR.